{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "thorough-brake",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "import pandas as pd\n",
    "from IPython.display import clear_output\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "\n",
    "from FBSM_library import FBSM, Objective_Functional\n",
    "from DM_library import DM\n",
    "from InitProblem_library import RKM_Dyn_System\n",
    "from Support_library import Define_Arguments, Obj_Functional_Iterations, Show_Control, Plot_Dynamics\n",
    "from Transition_Matrices_Bank import P_ass_strict_3, P_ass_strict_confidence_3, P_ass_smooth_3\n",
    "from Transition_Matrices_Bank import P_stubborn_3\n",
    "from Transition_Matrices_Bank import P_ass_smooth_5, P_BC_smooth_confidence_5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "146356ac-07c7-4343-adae-eca7f2c875a5",
   "metadata": {},
   "source": [
    "# m=3, M=2 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca169465-5190-4dad-bed6-401217a9be62",
   "metadata": {},
   "source": [
    "## Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43639ba0-7643-4c13-be42-f5976680bb98",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = 2 # the number of types of ordinary agents \n",
    "\n",
    "m = 3  # the number of elements in the opinion alphabet\n",
    "n = 0.6  # the fraction of ordinary agents in the system  \n",
    "\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "Transition_Matrices = {}\n",
    "\n",
    "# Initialize the first transition matrix\n",
    "\n",
    "# These transition matrices outline how ordinary agents influence each other \n",
    "\n",
    "Transition_Matrices['0-0'] = P_ass_strict_3\n",
    "\n",
    "Transition_Matrices['0-1'] = P_ass_strict_3\n",
    "\n",
    "Transition_Matrices['1-0'] = P_ass_strict_confidence_3\n",
    "\n",
    "Transition_Matrices['1-1'] = P_ass_strict_3\n",
    "\n",
    "# These transition matrices outline how stubborn agents influence ordinary agents \n",
    "\n",
    "Transition_Matrices['0-2'] = P_ass_strict_3\n",
    "\n",
    "Transition_Matrices['1-2'] = P_ass_strict_3\n",
    "    \n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the grid\n",
    "\n",
    "Tau_0 = 0\n",
    "Tau_1 = 40\n",
    "Step = 0.05\n",
    "Grid = np.arange(Tau_0, Tau_1+Step, Step)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Define the objective functional\n",
    "\n",
    "K = 1\n",
    "\n",
    "w = K*np.array([2, 1, 0])  # the opinion weights (integral)\n",
    "v = np.array([2, 1, 0])  # the opinion weights (terminal)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting point ( y_{0} )\n",
    "\n",
    "y_0 = np.array([[0.5,  0.1], \n",
    "                [0,    0], \n",
    "                [0,    0]])\n",
    "\n",
    "n_types = [0.5, 0.1]\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Create a dictionary of arguments\n",
    "\n",
    "Arguments = Define_Arguments(M, m, \n",
    "                             n, \n",
    "                             Transition_Matrices, \n",
    "                             Step, Grid, Tau_0, Tau_1, \n",
    "                             w, v, \n",
    "                             n_types)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting control guess ( u_{0} (\\tau) )\n",
    "\n",
    "u_basic = np.array([[0.2,  0.2], \n",
    "                    [0,    0], \n",
    "                    [0,    0]])\n",
    "\n",
    "u_t_0 = []\n",
    "\n",
    "for i in range(len(Grid)-1):\n",
    "    \n",
    "    u_t_0.append(u_basic) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f947c9dd-9655-40f5-b341-6cc27561d02a",
   "metadata": {},
   "source": [
    "## Run the direct method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91c747f8-10c1-44a1-a748-99125b2a5ca4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "u_t = DM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5164dce-cfdd-48ba-b921-01ebde6e13a9",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c5fc09e-4747-4070-9cb4-4726b870ff9d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Show the control function obtained\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The control function obtained:')\n",
    "\n",
    "print(Show_Control(u_t))\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Calculate the value of the objective functional\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The value of the objective functional provided by this control:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "510aa9dd-d759-42ae-ab4e-6703815c348c",
   "metadata": {},
   "source": [
    "## Run the FBS method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec414353-e292-4274-8bc2-877766edda4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "(u_t_estimations, Functional_estimations, Status) = FBSM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6056941a-01c1-434e-a19f-11c4a96a1daf",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0e8bbf8-ce6c-4516-815e-b8c1cf9d40c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The last approximation:')\n",
    "\n",
    "Show_Control(u_t_estimations[-1])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "print('')\n",
    "print('The best approximation:')\n",
    "\n",
    "sums = []\n",
    "\n",
    "for i in Functional_estimations:\n",
    "    sums.append(i[0]+i[1])\n",
    "\n",
    "Show_Control(u_t_estimations[np.argmin(sums)])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  # last approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('The value of the objective functional provided by the last approximation:', round(value[0]+value[1], 3))\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  # the best approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('')\n",
    "print('The value of the objective functional provided by the best approximation:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "199f8e3e-22ff-41a6-9f61-4f3f9bd812f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the dynamics of the objective functional across iterations of the FBS method\n",
    "\n",
    "FS = 15\n",
    "LW = 2\n",
    "MS = 5\n",
    "MEW = 2\n",
    "LS = 10\n",
    "\n",
    "Obj_Functional_Iterations(Functional_estimations, FS, LW, MS, MEW, LS)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the last approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the best approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8c574f5-8994-4053-a45c-e2bf006728a8",
   "metadata": {},
   "source": [
    "# m=5, M=2 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b58f7e-ac43-4ace-8718-4805f64f32ba",
   "metadata": {},
   "source": [
    "## Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53534cc5-e06e-4de2-99c2-3b57556c084c",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = 2 # the number of types of ordinary agents \n",
    "\n",
    "m = 5  # the number of elements in the opinion alphabet\n",
    "n = 0.6  # the fraction of ordinary agents in the system  \n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "Transition_Matrices = {}\n",
    "\n",
    "# Initialize the first transition matrix\n",
    "\n",
    "# These transition matrices outline how ordinary agents influence each other \n",
    "\n",
    "Transition_Matrices['0-0'] = P_ass_smooth_5\n",
    "\n",
    "Transition_Matrices['0-1'] = P_ass_smooth_5\n",
    "\n",
    "Transition_Matrices['1-0'] = P_BC_smooth_confidence_5\n",
    "\n",
    "Transition_Matrices['1-1'] = P_BC_smooth_confidence_5\n",
    "\n",
    "# These transition matrices outline how stubborn agents influence ordinary agents \n",
    "\n",
    "Transition_Matrices['0-2'] = P_ass_smooth_5\n",
    "\n",
    "Transition_Matrices['1-2'] = P_BC_smooth_confidence_5\n",
    "    \n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the grid\n",
    "\n",
    "Tau_0 = 0\n",
    "Tau_1 = 40\n",
    "Step = 1\n",
    "Grid = np.arange(Tau_0, Tau_1+Step, Step)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Define the objective functional\n",
    "\n",
    "K = 1\n",
    "\n",
    "w = K*np.array([4, 3, 2, 1, 0])  # the opinion weights (integral)\n",
    "v = np.array([4, 3, 2, 1, 0])  # the opinion weights (terminal)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting point ( y_{0} )\n",
    "\n",
    "y_0 = np.array([[0.1,  0.5], \n",
    "                [0,    0], \n",
    "                [0,    0], \n",
    "                [0,    0], \n",
    "                [0,    0]])\n",
    "\n",
    "n_types = [0.1, 0.5]\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Create a dictionary of arguments\n",
    "\n",
    "Arguments = Define_Arguments(M, m, \n",
    "                             n, \n",
    "                             Transition_Matrices, \n",
    "                             Step, Grid, Tau_0, Tau_1, \n",
    "                             w, v, \n",
    "                             n_types)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting control guess ( u_{0} (\\tau) )\n",
    "\n",
    "u_basic = np.array([[0.2,  0.2], \n",
    "                    [0,    0], \n",
    "                    [0,    0], \n",
    "                    [0,    0], \n",
    "                    [0,    0]])\n",
    "\n",
    "u_t_0 = []\n",
    "\n",
    "for i in range(len(Grid)-1):\n",
    "    \n",
    "    u_t_0.append(u_basic) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7df2858d-aee2-4e5c-a8c7-6ce4ec690da6",
   "metadata": {},
   "source": [
    "## Run the direct method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80a5e14c-8c83-4955-9985-fe1b9a104f0c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "u_t = DM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cdb4856-0589-47a1-a378-87f21c7126c1",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ea1cd1-b9ea-4e76-a163-b72343d134d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the control function obtained\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The control function obtained:')\n",
    "\n",
    "print(Show_Control(u_t))\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Calculate the value of the objective functional\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The value of the objective functional provided by this control:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d19e9e2-11d9-4512-b505-747894a3e7ea",
   "metadata": {},
   "source": [
    "## Run the FBS method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7d25ffb-b64b-411f-9ffa-cf44d4851fa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "(u_t_estimations, Functional_estimations, Status) = FBSM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1803427-edbd-43ba-a51f-12f452ebea9d",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6637e44-4f44-44ec-add7-53c73562cef9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The last approximation:')\n",
    "\n",
    "Show_Control(u_t_estimations[-1])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "print('')\n",
    "print('The best approximation:')\n",
    "\n",
    "sums = []\n",
    "\n",
    "for i in Functional_estimations:\n",
    "    sums.append(i[0]+i[1])\n",
    "\n",
    "Show_Control(u_t_estimations[np.argmin(sums)])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  # last approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('The value of the objective functional provided by the last approximation:', round(value[0]+value[1], 3))\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  # the best approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('')\n",
    "print('The value of the objective functional provided by the best approximation:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fd50717-33cd-4378-ab0e-329e1d4bcaeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the dynamics of the objective functional across iterations of the FBS method\n",
    "\n",
    "FS = 15\n",
    "LW = 2\n",
    "MS = 5\n",
    "MEW = 2\n",
    "LS = 10\n",
    "\n",
    "Obj_Functional_Iterations(Functional_estimations, FS, LW, MS, MEW, LS)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the last approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the best approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
